{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TensorRT Engine Explorer Tutorial\n",
    "\n",
    "Use this notebook to learn how to use `trex` to explore the structure and characteristics of a TensorRT Engine plan.\n",
    "Starting with TensorRT 8.2, engine-plan graph and profiling data can be exported to JSON files.  `trex` loads these files and queries their content using a simple API that wraps Pandas' API"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Create Engine Plan And Profiling JSON Files\n",
    "\n",
    "Using `trtexec` it is easy to create the plan graph and profiling JSON files.  Three flags are required:\n",
    "```\n",
    " --exportProfile=$profile_json\n",
    " --exportLayerInfo=$graph_json\n",
    " --profilingVerbosity=detailed\n",
    "```\n",
    "\n",
    "A utility Python script `utils/process_engine.py` can be used to create the JSON files. This script executes `trtexec` and therefore the script can only be invoked within an environment that has `trtexec` installed and included in `$PATH`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !python3 ../utils/process_engine.py ../tests/inputs/mobilenet.qat.onnx ../tests/inputs best"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Load JSON Files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "import pandas as pd\n",
    "import trex\n",
    "import trex.notebook\n",
    "import trex.plotting\n",
    "import trex.graphing\n",
    "import trex.df_preprocessing\n",
    "\n",
    "# Configure a wider output (for the wide graphs)\n",
    "trex.notebook.set_wide_display()\n",
    "\n",
    "# Choose an engine file to load.  This notebook assumes that you've saved the engine to the following paths.\n",
    "engine_name = \"../tests/inputs/mobilenet.qat.onnx.engine\"\n",
    "engine_name = \"../tests/inputs/mobilenet_v2_residuals.qat.onnx.engine\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Create an `EnginePlan` instance and start exploring"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert engine_name is not None\n",
    "plan = trex.EnginePlan(f'{engine_name}.graph.json', f'{engine_name}.profile.json', f'{engine_name}.profile.metadata.json')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "It is helpful to look at a high-level summary of the engine plan before diving into the details.\n",
    "\n",
    "* By default an `EnginePlan`'s name is derived from the graph file name, but you can also set the name explicitly at any time. This is useful when there are multiple plans with the same file name.\n",
    "* The data is summarized from the input JSON files, and may include results from profiling layers separately as well as from profiling the entire engine. The latter provides more accurate values for latency and inferences throughput, but the former provides per-layer latency information for understanding the engine behavior.\n",
    "* \"Average time\": refers to the sum of the layer latencies, when profiling layers separately.\n",
    "* \"Latency\": refers to the [min, max, mean, median, 99% percentile] of the engine latency measurements, when timing the engine w/o profiling layers.\n",
    "* \"Throughput\": is measured in inferences per second (IPS)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Summary for {plan.name}:\\n\")\n",
    "plan.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get to know the plan dataframe\n",
    "\n",
    "An `EnginePlan` is an object that wraps a Pandas DataFrame data-structure.  Most of the examples below utilize this dataframe (df) for querying, slicing and rendering information about the `EnginePlan`.\n",
    "\n",
    "The dataframe captures the information from the plan engine graph and profiling JSON files.  If both JSON files are available the latency data of each layer is added as three new columns: `latency.time` (the total latency of the layer summed across all measurement iterations), `latency.avg_time` (the average latency of the layer), `latency.pct_time` (the latency of the layer as a proportion of the overall engine latency).\n",
    "\n",
    "When the dataframe is constructed several new layers that summarize footprint information are added: `total_io_size_bytes`, `weights_size`, `total_footprint_bytes`.\n",
    "\n",
    "A few of the column names are changed from the original JSON file, to give them clearer names.\n",
    "\n",
    "Accessing the dataframe is straight-forward:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = plan.df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can print the names of the columns in the dataframe."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "available_cols = df.columns\n",
    "print(f\"These are the column names in the plan\\n: {available_cols}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A dataframe can be rendered as a table.  The columns are from various layers so the dataframe is very sparse.\n",
    "\n",
    "Use the column controls to sort or filter layers.<br>\n",
    "An interesting view sorts the layers by `latency.pct_time`.<br>\n",
    "The `dtale` toolbar makes it easy to open the table in a new tab (useful for large tables) and to export the data to CSV and HTML."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trex.notebook.display_df(plan.df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When rendering engine plan dataframes we usually want to reduce the visual clutter and render only the important columns.<br>\n",
    "The function `clean_for_display` does exactly that.\n",
    "\n",
    "The column order is changed, in order to bring important columns to the front.<br>\n",
    "Columns `Inputs` and `Outputs` are reformatted to reduce verbosity.<br>\n",
    "Finally, a few columns are dropped and NaNs are replaced with zeros."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = trex.df_preprocessing.clean_for_display(plan.df)\n",
    "print(f\"These are the column names in the plan\\n: {df.columns}\")\n",
    "trex.notebook.display_df(df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Layer Types\n",
    "\n",
    "This example shows how to create a bar diagram of the count of each layer type.\n",
    "\n",
    "`trex` provides a utility wrapper around Pandas' API, but you can freely use the Pandas API to extract data from the plan dataframe."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "layer_types = trex.group_count(plan.df, 'type')\n",
    "\n",
    "# Simple DF print\n",
    "print(layer_types)\n",
    "\n",
    "# dtale DF display\n",
    "trex.notebook.display_df(layer_types)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`trex` provides wrappers to plotly's plotting API.  `plotly_bar2` is the main utility for creating bar charts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trex.plotting.plotly_bar2(\n",
    "    df=layer_types, \n",
    "    title='Layer Count By Type', \n",
    "    values_col='count', \n",
    "    names_col='type',\n",
    "    orientation='v',\n",
    "    color='type',\n",
    "    colormap=trex.colors.layer_colormap,\n",
    "    show_axis_ticks=(True, True));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Performance\n",
    "\n",
    "Pandas' powerful API can be used on the Plan dataframe.  For example, we can easily query for the 3 layers that consume the most time:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "top3 = plan.df.nlargest(3, 'latency.pct_time')\n",
    "trex.notebook.display_df(top3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The chart below provides a quick view of the layers latencies. The values of `df[values_col]` set the bar height and the values of `df[names_col]` provide the bar name. In this case, the latency of each layer is plotted vs the name of the layer. The colors of the bars are determined by `color` and `colormap`, if provided.<br>\n",
    "For example in the statement `colormap[df['type']]`, the bar colors are determined by the layer type and `layer_colormap`, which is a `trex` dictionary which maps layer types to preset colors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trex.plotting.plotly_bar2(\n",
    "    df=plan.df, \n",
    "    title=\"% Latency Budget Per Layer\",\n",
    "    values_col=\"latency.pct_time\",\n",
    "    names_col=\"Name\",\n",
    "    color='type',\n",
    "    use_slider=False,\n",
    "    colormap=trex.colors.layer_colormap);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`plotly_hist` is a wrapper of Plotly's histograms chart. It has arguments similar to `plotly_bar2`, but not as many. `plotly_hist` plots the histogram of `df[values_col]`.\n",
    "\n",
    "Here's a look at how the layer latencies distribute:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trex.plotting.plotly_hist(\n",
    "    df=plan.df, \n",
    "    title=\"Layer Latency Distribution\", \n",
    "    values_col=\"latency.pct_time\",\n",
    "    xaxis_title=\"Latency (ms)\",\n",
    "    color='type',\n",
    "    colormap=trex.colors.layer_colormap);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Pandas' aggregation and reductions can be used to provide interesting information.\n",
    "\n",
    "Here we group the layer latencies by the layer types.<br>The data can be displayed as a chart or as a summary table, like the one below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "time_pct_by_type = plan.df.groupby([\"type\"])[[\"latency.pct_time\", \"latency.avg_time\"]].sum().reset_index()\n",
    "trex.notebook.display_df(time_pct_by_type)\n",
    "trex.plotting.plotly_bar2(\n",
    "    df=time_pct_by_type,\n",
    "    title=\"% Latency Budget Per Layer Type\",\n",
    "    values_col=\"latency.pct_time\",\n",
    "    names_col=\"type\",\n",
    "    orientation='h',\n",
    "    color='type',\n",
    "    colormap=trex.colors.layer_colormap);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Treemaps provide a different view of the profiling data.\n",
    "\n",
    "In this example we use a Plotly Express Treemap directly, without any wrappers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import plotly.express as px\n",
    "\n",
    "fig = px.treemap(\n",
    "    plan.df,\n",
    "    path=['type', 'Name'],\n",
    "    values='latency.pct_time',\n",
    "    title='Treemap Of Layer Latencies (Size & Color Indicate Latency)',\n",
    "    color='latency.pct_time')\n",
    "fig.update_layout(margin = dict(t=50, l=25, r=25, b=25))\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is another view of how layer latencies interact with layer data size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = px.treemap(\n",
    "    plan.df,\n",
    "    path=['type', 'Name'],\n",
    "    values='latency.pct_time',\n",
    "    title='Treemap Of Layer Latencies (Size Indicates Latency. Color Indicates Activations Size)',\n",
    "    color='total_io_size_bytes')\n",
    "fig.update_traces(root_color=\"white\")\n",
    "fig.update_layout(margin = dict(t=50, l=25, r=25, b=25))\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Memory Traffic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trex.plotting.plotly_bar2(\n",
    "    plan.df, \n",
    "    \"Weights Sizes Per Layer\", \n",
    "    \"weights_size\", \"Name\", \n",
    "    color='type', \n",
    "    colormap=trex.colors.layer_colormap)\n",
    "\n",
    "trex.plotting.plotly_bar2(\n",
    "    plan.df, \n",
    "    \"Activations Sizes Per Layer\", \n",
    "    \"total_io_size_bytes\", \n",
    "    \"Name\", \n",
    "    color='type', \n",
    "    colormap=trex.colors.layer_colormap)\n",
    "\n",
    "trex.plotting.plotly_hist(\n",
    "    plan.df, \n",
    "    \"Layer Activations Sizes Distribution\", \n",
    "    \"total_io_size_bytes\", \n",
    "    \"Size (bytes)\", \n",
    "    color='type', \n",
    "    colormap=trex.colors.layer_colormap)\n",
    "\n",
    "plan.df[\"total_io_size_bytes\"].describe()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Layers Precision\n",
    "\n",
    "`trex` provides a wrapper to Plotly's pie charts as well. Several charts can be plotted in a grid.\n",
    "\n",
    "`precision_colormap` color the pie slices by the value of `df['precision']`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "charts = []\n",
    "layer_precisions = trex.group_count(plan.df, 'precision')\n",
    "charts.append((layer_precisions, 'Layer Count By Precision', 'count', 'precision'))\n",
    "\n",
    "layers_time_pct_by_precision = trex.group_sum_attr(plan.df, grouping_attr='precision', reduced_attr='latency.pct_time')\n",
    "display(layers_time_pct_by_precision)\n",
    "\n",
    "charts.append((layers_time_pct_by_precision, '% Latency Budget By Precision', 'latency.pct_time', 'precision'))\n",
    "trex.plotting.plotly_pie2(\"Precision Statistics\", charts, colormap=trex.colors.precision_colormap);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trex.plotting.plotly_bar2(\n",
    "    plan.df, \n",
    "    \"% Latency Budget Per Layer<BR>(bar color indicates precision)\", \n",
    "    \"latency.pct_time\", \n",
    "    \"Name\",\n",
    "    color='precision',\n",
    "    colormap=trex.colors.precision_colormap);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Graph Rendering\n",
    "\n",
    "It is very helpful to draw the graph of the engine plan.\n",
    "\n",
    "A formatter can be used to configure the colors of nodes. `trex` provides `layer_type_formatter` which paints graph nodes by their layer type, and `precision_formatter` which paints graph nodes according to their precision.\n",
    "\n",
    "`to_dot` converts an `EnginePlan` to dot file which can be rendered to SVG or PNG.\n",
    "\n",
    "SVG files render faster than PNG, they are searchable and provide sharp and crisp graphs in all resolutions.  Because graphs are large, it is recommended to view the rendered graph file in another browser window."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "formatter = trex.graphing.layer_type_formatter if True else trex.graphing.precision_formatter\n",
    "graph = trex.graphing.to_dot(plan, formatter)\n",
    "svg_name = trex.graphing.render_dot(graph, engine_name, 'svg')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "PNG files can be rendered inside the notebook, but the graphs are usually very large and resolution suffers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "png_name = trex.graphing.render_dot(graph, engine_name, 'png')\n",
    "from IPython.display import Image\n",
    "display(Image(filename=png_name))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convolution Layers\n",
    "\n",
    "Sometimes it is interesting to look at all layers of a certain type. You can use Pandas' API (e.g. `query`) to slice the dataframe by layer type:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "convs1 = plan.df.query(\"type == 'Convolution'\")\n",
    "convs2 = df[df.type == 'Convolution']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, `trex` provides a `get_layers_by_type` API which performs some layer-type-specific preprocessing which is often useful. In the case of convolutions,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "convs = plan.get_layers_by_type('Convolution')\n",
    "trex.notebook.display_df(convs)\n",
    "\n",
    "trex.plotting.plotly_bar2(\n",
    "    convs, \n",
    "    \"Latency Per Layer (%)<BR>(bar color indicates precision)\",\n",
    "    \"attr.arithmetic_intensity\", \"Name\",\n",
    "    color='precision', \n",
    "    colormap=trex.colors.precision_colormap)\n",
    "\n",
    "trex.plotting.plotly_bar2(\n",
    "    convs,\n",
    "    \"Convolution Data Sizes<BR>(bar color indicates latency)\",\n",
    "    \"total_io_size_bytes\", \n",
    "    \"Name\", \n",
    "    color='latency.pct_time');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Arithmetic Intensity\n",
    "\n",
    "Arithmetic intensity (AI) is a measure of the amount of compute expended per data byte. Layers with higher AI are in general more efficient because moving data is much slower than computing an operation.  For each unit of data fetched from memory we want to perform many computations, because the GPU can compute much faster than it can fetch data from memory.  From more on AI, see https://en.wikipedia.org/wiki/Roofline_model#Arithmetic_intensity.<br>\n",
    "This is a simplistic model which assumes that we read the data only once, and write it out once. In practice, when computing Convolution and GEMM operations, memory tiles are read several times, usually from fast shared-memory, L1 or L2 cached memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trex.plotting.plotly_bar2(\n",
    "    convs, \n",
    "    \"Convolution Arithmetic Intensity<BR>(bar color indicates activations size)\",\n",
    "    \"attr.arithmetic_intensity\", \n",
    "    \"Name\",\n",
    "    color='total_io_size_bytes')\n",
    "\n",
    "trex.plotting.plotly_bar2(\n",
    "    convs, \n",
    "    \"Convolution Arithmetic Intensity<BR>(bar color indicates latency)\", \n",
    "    \"attr.arithmetic_intensity\", \n",
    "    \"Name\",\n",
    "    color='latency.pct_time');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Efficiency\n",
    "\n",
    "Another simplistic model measures the compute and memory efficiency.  These indicators are calculated by taking the number of operations (or memory bytes) and dividing by the layer's execution time. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Memory accesses per ms (assuming one time read/write penalty)\n",
    "trex.plotting.plotly_bar2(\n",
    "    convs, \n",
    "    \"Convolution Memory Efficiency<BR>(bar color indicates latency)\", \n",
    "    \"attr.memory_efficiency\", \n",
    "    \"Name\", \n",
    "    color='latency.pct_time')\n",
    "\n",
    "# Compute operations per ms (assuming one time read/write penalty)\n",
    "trex.plotting.plotly_bar2(\n",
    "    convs, \n",
    "    \"Convolution Compute Efficiency<BR>(bar color indicates latency)\",\n",
    "    \"attr.compute_efficiency\",\n",
    "    \"Name\",\n",
    "    color='latency.pct_time');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Statistics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "convs = plan.get_layers_by_type('Convolution')\n",
    "\n",
    "charts = []\n",
    "convs_count_by_type = trex.group_count(convs, 'subtype')\n",
    "charts.append((convs_count_by_type, 'Count', 'count', 'subtype'))\n",
    "\n",
    "convs_time_pct_by_type = trex.group_sum_attr(convs, grouping_attr='subtype', reduced_attr='latency.pct_time')\n",
    "charts.append((convs_time_pct_by_type, '% Latency Budget', 'latency.pct_time', 'subtype'))\n",
    "trex.plotting.plotly_pie2(\"Convolutions Statistics (Subtype)\", charts)\n",
    "\n",
    "charts = []\n",
    "convs_count_by_group_size = trex.group_count(convs, 'attr.groups')\n",
    "charts.append((convs_count_by_group_size, 'Count', 'count', 'attr.groups'))\n",
    "\n",
    "convs_time_pct_by_grp_size = trex.group_sum_attr(convs, grouping_attr='attr.groups', reduced_attr='latency.pct_time')\n",
    "charts.append((convs_time_pct_by_grp_size, '% Latency Budget', 'latency.pct_time', 'attr.groups'))\n",
    "trex.plotting.plotly_pie2(\"Convolutions Statistics (Number of Groups)\", charts)\n",
    "\n",
    "charts = []\n",
    "convs_count_by_precision = trex.group_count(convs, 'precision')\n",
    "charts.append((convs_count_by_precision, 'Count', 'count', 'precision'))\n",
    "\n",
    "convs_time_pct_by_precision = trex.group_sum_attr(convs, grouping_attr='precision', reduced_attr='latency.pct_time')\n",
    "charts.append((convs_time_pct_by_precision, '% Latency Budget', 'latency.pct_time', 'precision'))\n",
    "\n",
    "trex.plotting.plotly_pie2(\"Convolutions Statistics (Precision)\", charts, colormap=trex.colors.precision_colormap);"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
